from itertools import chain
import os
from typing import Sequence
from feasytools import TimeImplictFunc
from fpowerkit.grid import Grid
from fstatistics import *
from .evloadpred import *
from fplugins import *


class PluginDynPsWithAI(PluginBase):
    @property
    def Description(self)->str:
        return "动态售电价格机制"
        
    def Initialization(self,elem:ET.Element,inst:TrafficInst,work_dir:Path,plugin_dependency:'list[PluginBase]') -> None:
        # 加载电网
        grid_file = elem.attrib.pop("grid")
        assert grid_file is not None, "未指定电网文件"
        grid_path = Path(grid_file)
        if grid_path.is_absolute(): grid_file = str(grid_path)
        else: grid_file = str(work_dir / grid_path)
        self.__gr = Grid.fromFile(grid_file, True)
        self.__cur_pd_pu:Sequence[float] = [0.] * (len(inst.CSList) + len(inst.PKList))

        # 将各充电站注册至母线负荷
        for i, c in enumerate(chain(inst.CSList,inst.PKList)):
            if not c.node in self.__gr.BusNames: raise ValueError(f"错误: 充电站{c.name}的母线{c.node}在电网中不存在")
            self.__gr.Bus(c.node).Pd += TimeImplictFunc(lambda: self.__cur_pd_pu[i])
        
        # 加载GEN-CS绑定
        slot_change = elem.find("slotschange")
        self.__slotsch:list[tuple[str,list[str]]] = []
        if slot_change:
            for e in slot_change:
                self.__slotsch.append((e.attrib.pop("gen"),e.attrib.pop("cs").split(",")))

        # 加载调价参数
        self.__inst = inst
        attr = elem.attrib
        self.__offset = float(attr.pop("offset",0))
        self.__future = int(attr.pop("future"))
        self.__sr = int(attr.get("sr",1))
        assert self.__future % self.__sr == 0, "未来预测时间应该是采样率的整数倍"

        # 加载AI预测模型
        os.chdir("./external_plugins/evloadpred")
        #scaler_path = "./ckpts/"+attr["data"]+"/scaler.pkl"
        #with open(scaler_path,"rb") as f:
        #    self.__scaler = pickle.load(f)
        server_addr = ("localhost", 9999)
        self.__pred = TrafficPredictorClient(attr,server_addr)
        os.chdir("../../")
        self.__pred.set_trips(ET.ElementTree(file = work_dir / (work_dir.name + ".veh.xml")).getroot())
        self.SetPreStep(self.Work)
        self.__can_pred = False
        return None

    def Work(self,t:int,sta:PluginStatus)->tuple[bool,None]:
        # 无论是否运行，都应该记录数据
        if t >= 86400:
            # AI预测使用的数据单位是kW
            self.__can_pred = self.__pred.add_data(t - 86400, chain(
                self.__inst.CSList.get_prices_at(t),
                self.__inst.PKList.get_prices_at(t),
                self.__inst.CSList.get_Pc(3600),
                self.__inst.PKList.get_Pc(3600)
            ))

        if sta != PluginStatus.EXECUTE: return True, None
        # 如果执行，应该进行预测
        if self.__can_pred:
            sb_kVA = self.__gr.Sb_kVA
            tl, dt = self.__pred.predict()
            # AI预测使用的数据单位是kW, 要转换成标幺值
            self.__cur_pd_pu = list(float(x)/sb_kVA for x in dt[self.__future // self.__sr - 1,:])
            #print(f"Pred: {tl[0]}, {tl[-1]}")
        else:
            k = 3600/self.__gr.Sb_kVA
            self.__cur_pd_pu = self.__inst.CSList.get_Pc(k) + self.__inst.PKList.get_Pc(k)

        if self.__can_pred:
            ret, val = self.__gr.solve(t + self.__future * 60)
        else:
            ret, val = self.__gr.solve(t)
        if ret:
            for cs in self.__inst.CSList:
                cshd = self.__gr.Bus(cs.node).ShadowPrice
                if cshd is None:
                    cs.pbuy.clearOverride()
                else:
                    cs.pbuy.setOverride(cshd / self.__gr.Sb_kVA + self.__offset)
            return True, None
        else:
            return False, None
    
exports = (PluginDynPsWithAI, [])